{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\nExternal Tensor Functions\n=========================\n**Author**: `Tianqi Chen <https://tqchen.github.io>`_\n\nWhile TVM supports transparent code generation, sometimes\nit is also helpful to incorporate manual written code into\nthe pipeline. For example, we might want to use cuDNN for\nsome of the convolution kernels and define the rest of the stages.\n\nTVM supports these black box function calls natively.\nSpecfically, TVM support all the tensor functions that are DLPack compatible.\nWhich means we can call any function with POD types(pointer, int, float)\nor pointer to DLTensor as argument.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from __future__ import absolute_import, print_function\n\nimport tvm\nfrom tvm import te\nimport numpy as np\nfrom tvm.contrib import cblas\nimport tvm.testing\n\nif not tvm.get_global_func(\"tvm.contrib.cblas.matmul\", allow_missing=True):\n    raise Exception(\"Not compiled with cblas support; can't build this tutorial\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Use Extern Tensor Function\n--------------------------\nIn the example below, we use :any:`te.extern` to add an extern\narray function call. In the extern call, we declare the shape\nof output tensors. In the second argument we provide the list of inputs.\n\nUser will need to provide a function describing how to compute the result.\nThe compute function takes list of symbolic placeholder for the inputs,\nlist of symbolic placeholder for the outputs and returns the executing statement.\n\nIn this case we simply call a registered TVM function, which invokes a CBLAS call.\nTVM does not control internal of the extern array function and treats it as blackbox.\nWe can further mix schedulable TVM calls that add a bias term to the result.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "n = 1024\nl = 128\nm = 235\nbias = te.var(\"bias\", dtype=\"float32\")\nA = te.placeholder((n, l), name=\"A\")\nB = te.placeholder((l, m), name=\"B\")\nC = te.extern(\n    (n, m),\n    [A, B],\n    lambda ins, outs: tvm.tir.call_packed(\n        \"tvm.contrib.cblas.matmul\", ins[0], ins[1], outs[0], False, False\n    ),\n    name=\"C\",\n)\nD = te.compute(C.shape, lambda i, j: C[i, j] + bias, name=\"D\")\ns = te.create_schedule(D.op)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Verify the Result\n-----------------\nWe can verify that the result matches what we expected.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "dev = tvm.cpu(0)\nf = tvm.build(s, [A, B, D, bias], \"llvm\")\na = tvm.nd.array(np.random.uniform(size=(n, l)).astype(A.dtype), dev)\nb = tvm.nd.array(np.random.uniform(size=(l, m)).astype(B.dtype), dev)\nd = tvm.nd.array(np.zeros((n, m), dtype=D.dtype), dev)\nbb = 10.0\nf(a, b, d, bb)\ntvm.testing.assert_allclose(d.numpy(), np.dot(a.numpy(), b.numpy()) + 10, rtol=1e-5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Extern Contrib Wrappers\n-----------------------\nTVM also provide extern contrib wrappers to useful extern calls,\nthe following line is equivalent to the previous example.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from tvm.contrib import cblas\n\nC = cblas.matmul(A, B)\nD = te.compute(C.shape, lambda i, j: C[i, j] + bias, name=\"D\")\ns = te.create_schedule(D.op)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Hook Python Function as Extern\n------------------------------\nSince we can call into any PackedFunc in TVM. We can use the extern\nfunction to callback into python.\n\nThe following example registers a python function into TVM runtime system\nand use it to complete one stage of the computation.\nThis makes TVM much more flexible. For example, we can insert front-end\ncallbacks to inspect the intermediate results or mix customized code\nwith TVM.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "@tvm.register_func(\"tvm.contrib.my_tvm_addone\")\ndef my_tvm_addone(x, y):\n    print(\"my_tvm_addone signatures: %s, %s\" % (type(x), type(y)))\n    tvm.nd.array(x.numpy() + 1).copyto(y)\n\n\nA = te.placeholder((n,), name=\"A\")\nB = te.extern(\n    A.shape,\n    [A],\n    lambda ins, outs: tvm.tir.call_packed(\"tvm.contrib.my_tvm_addone\", ins[0], outs[0]),\n    name=\"C\",\n)\ns = te.create_schedule(B.op)\nf = tvm.build(s, [A, B], \"llvm\")\na = tvm.nd.array(np.random.uniform(size=(n,)).astype(A.dtype), dev)\nb = tvm.nd.array(np.random.uniform(size=(n,)).astype(B.dtype), dev)\nf(a, b)\ntvm.testing.assert_allclose(b.numpy(), a.numpy() + 1, rtol=1e-5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Summary\n-------\n- TVM calls extern tensor function via :any:`te.extern`\n- Use contrib wrappers for short sugars of extern tensor calls.\n- We can hook front-end function as extern tensor callbacks.\n\n\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}